# 用一个类定义一个节点
from typing import Any

from langchain_core.messages import HumanMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import Graph,END

from rich.console import Console
console = Console()

from rich.markdown import Markdown

# 初始化模型
def qw_model():
    return ChatOpenAI(
        model='qwen-max-0919',
        base_url='https://dashscope.aliyuncs.com/compatible-mode/v1',
        api_key='sk-debee146f82244268914dd2e3d98761b',
        temperature=0.7,
        top_p=0.9
    )

model = qw_model()

def entry(input:list[HumanMessage]):
    return input

def action(input:list[HumanMessage]):
    print("Action taken:", [msg.content for msg in input])
    if len(input) > 5:
        input.append(HumanMessage(content="end"))
    else:
        input.append(HumanMessage(content="continue"))
    return input

def condition(input:list):
    last_message = input[-1]
    if "end" in last_message.content:
        return "__end__"
    return "action"


# 2.构建图
workflow = Graph()
# 添加节点
workflow.add_node("entry", entry)
workflow.add_node("action", action)
# 添加条件边
workflow.add_conditional_edges("entry",condition,{"action": "action", "__end__":END})

workflow.add_edge("action", "entry")

# 定义入口节点和出口节点
workflow.set_entry_point("entry")

# 3.对图进行编译
app = workflow.compile()
# 调用方式
res= app.invoke([HumanMessage(content="hello")])
print(res)